/*
 * EquilibriumSolver.cpp
 *
 *  Created on: 26 Jul 2011
 *      Author: Allan
 */

#include "EquilibriumSolver.h"

// C++ includes
#include <cmath>
#include <functional>
#include <iomanip>
#include <memory>
using namespace std;
using namespace std::placeholders;

// Boost includes
#include <boost/foreach.hpp>
using namespace boost;

// GeoReact includes
#include "../SpeciesMap.h"
#include "../Utilities/Algorithms.h"
#include "../Utilities/SpeciesUtils.h"

EquilibriumSolver::EquilibriumSolver() :
tolerance(1.0E-6), maxiterations(100)
{}

EquilibriumSolver::EquilibriumSolver(const Multiphase& multiphase, const CanonicalReactionSystem& reactions) :
tolerance(1.0E-6), maxiterations(100)
{
	Initialiaze(multiphase, reactions);
}

void EquilibriumSolver::Initialiaze(const Multiphase& multiphase_, const CanonicalReactionSystem& reactions_)
{
	// Set the multiphase object of the class
	multiphase = multiphase_;
	
	// Set the equilibrium reactions in the canonical form
	reactions = reactions_;
	
	// Set the equilibrium problem object
	problem = EquilibriumProblem(multiphase, reactions);
	
	// Set the primary (jSpecies) and equilibrium (jSpecies) species
	jSpecies = reactions.GetPrimarySpecies();
	eSpecies = reactions.GetSecondarySpecies();
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Nj = jSpecies.size();
	Ne = eSpecies.size();
	
	// Set the indexes of the primary (jIndexes) and equilibrium (eIndexes) species
	jIndexes = multiphase[jSpecies];
	eIndexes = multiphase[eSpecies];
	
	// Initialiaze the multiroot solver
	solver.Initialiaze(Nj + Ne, StandardNewton);
	
	// Set the algorithm of the solver to the standard Newton-Raphson method 
	// Note that the modified Newton-Raphson algorithm works too, but the algorithms Hybridj and Hybridsj does not 
	solver.SetSolverOptions(tolerance, maxiterations);
	
	// Set the stoichiometric matrix of the system of reactions
	vej = reactions.GetCanonicalStoichiometricMatrix();
}

void EquilibriumSolver::SetEquilibriumConditions(const vector<EquilibriumCondition>& eConditions)
{
	problem.SetEquilibriumConditions(eConditions);
}

void EquilibriumSolver::SetSolverOptions(double tolerance, unsigned maxiterations)
{
	this->tolerance = tolerance;
	this->maxiterations = maxiterations;
	
	solver.SetSolverOptions(tolerance, maxiterations);
}

const EquilibriumSolver::Diagnostic EquilibriumSolver::Solve(double T, double P, VectorXd& n)
{
	// Set the paramenters of the equilibrium problem
	problem.SetParameters(T, P, n);
	
	// The vector of moles of primary/equilibrium species (njne) = [nj ne]
	VectorXd njne(Nj + Ne);
	
	// Initialise the vector (njne) with the values from the initial guess (nGuess)
	for(unsigned j = 0; j < Nj; ++j) njne[j]      = n[jIndexes[j]];
	for(unsigned e = 0; e < Ne; ++e) njne[e + Nj] = n[eIndexes[e]];
	
	// Set the diagnostic instance
	Diagnostic diagnostic = {0, false};
	
	// The solution vector (njne) at the previous iteration (used for application of under relaxation to the solution)
	VectorXd njneOld = njne;
	
	// Solve the equilibrium problem
	while(solver.Iterate(problem, njne) && diagnostic.iterations < maxiterations)
	{
		// Apply the under relaxation to the vector (njne) to avoid negative values of compositions
		UnderRelax(njne, njneOld);
		
		// Update the vector (njneOld)
		njneOld = njne;
		
		// Update (iteration)
		++diagnostic.iterations;
	}
	
	// Check convergence of the previous iterations
	diagnostic.converged = diagnostic.iterations < maxiterations;
	
	// Transfer the solution (njne) to the vector (n)
	for(unsigned j = 0; j < Nj; ++j) n[jIndexes[j]] = njne[j];
	for(unsigned e = 0; e < Ne; ++e) n[eIndexes[e]] = njne[e + Nj];
	
	return diagnostic;
}

void EquilibriumSolver::UnderRelax(VectorXd& n, const VectorXd& nOld) const
{
	for(unsigned i = 0; i < unsigned(n.rows()); ++i)
	{
		const double delta  = n[i] - nOld[i];
		const double factor = max(1.0, -2.0 * delta/nOld[i]);
		
		n[i] = nOld[i] + delta/factor;
	}
}

const EquilibriumCondition EquilibriumSolver::FixedMassHelper(const string& species, const FunctionTPna& h)
{
	// The equilibrium condition instance
	EquilibriumCondition eCondition(Nj, Ne);
	
	// Set the equilibrium condition function (h)
	eCondition.h = h;
	
	// Get the local indexes of the given (species) in the set of primary (jSpecies) and equilibrium (eSpecies) species
	const unsigned j = IndexOf(species, jSpecies);
	const unsigned e = IndexOf(species, eSpecies);
	
	// If the given (species) is a primary species, set the alphaRow vector of (eCondition)
	if(j < Nj) eCondition.alphaRow[j] = 1.0;
	
	// If the given (species) is an equilibrium species, set the betaRow vector of (eCondition)
	if(e < Ne) eCondition.betaRow[e] = 1.0;
	
	return eCondition;
}

const EquilibriumCondition EquilibriumSolver::FixedMass(const string& species, double mass)
{
	return FixedMassHelper(species, ConstantTPna(mass));
}

const EquilibriumCondition EquilibriumSolver::FixedTotalMass(const string& species, double totalmass)
{
	// The equilibrium condition instance
	EquilibriumCondition eCondition = FixedMass(species, totalmass);
	
	// Get the local index of the given (species) in the set of primary species (jSpecies)
	const unsigned j = IndexOf(species, jSpecies);
	
	// If the given (species) is a primary species, set the betaRow vector of (eCondition)
	if(j < Nj) eCondition.betaRow = vej.col(j);
	
	return eCondition;
}

const EquilibriumCondition EquilibriumSolver::FixedActivity(const string& species, double activity)
{
	// The index of the given species in the (multiphase) instance
	const Index i = multiphase[species];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	FunctionTPna h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
		{ return n[i]/a[i] * activity; };
	
	return FixedMassHelper(species, h);
}

const EquilibriumCondition EquilibriumSolver::FixedAcidity(double pH)
{
	return FixedActivity("H[+]", pow(10.0, -pH));
}

const EquilibriumCondition EquilibriumSolver::FixedPartialPressure(const string& gas, double Pg)
{
	// The phases in the (multiphase) instance
	const vector<Phase>& phases = multiphase.GetPhases();
	
	// Define a function that determines if a phase is gaseous
	auto isGaseous = [](const Phase& phase) { return phase.GetName() == "Gaseous"; };
	
	// Find the gaseous phase among all the phases in the system
	const Phase& gaseousPhase = *find_if(phases.begin(), phases.end(), isGaseous);
	
	// The indexes of all gaseous species
	vector<Index> igases = multiphase[gaseousPhase.GetSpecies()];
	
	// Define the right-hand side function h = h(T, P, n, a) for this equilibrium condition
	auto h = [=](double T, double P, const VectorXd& n, const VectorXd& a)
	{ 
		// The total number of moles of gaseous species
		double nTG = 0.0; for(unsigned i = 0; i < igases.size(); ++i) nTG += n[igases[i]];
		
		return nTG * Pg/P;
	};
	
	return FixedMassHelper(gas, h);
}

const EquilibriumCondition EquilibriumSolver::ChargeBalanced()
{
	// The equilibrium condition instance
	EquilibriumCondition eCondition(Nj, Ne);
	
	// Set the alphaRow vector of (eCondition) with the electrical charges of the primary species (jSpecies)
	for(unsigned j = 0; j < Nj; ++j) eCondition.alphaRow[j] = ElectricalCharge(jSpecies[j]);
	
	// Set the betaRow vector of (eCondition) with the electrical charges of the equilibrium species (eSpecies)
	for(unsigned e = 0; e < Ne; ++e)  eCondition.betaRow[e] = ElectricalCharge(eSpecies[e]);
	
	// Set the equilibrium condition function (h)
	eCondition.h = ConstantTPna(0.0);
	
	return eCondition;
}
